# -*- coding: utf-8 -*-
"""
Rerank模型对比实验示例
展示如何使用embedding+rerank和rerank only对比工具
"""

import os
import sys
import json
from typing import List
from loguru import logger

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 现在可以正常导入项目模块
from questionretrieval.config import *  # 设置HuggingFace镜像环境变量

from questionretrieval.retrievers.rerank.embedding_rerank_comparison import EmbeddingRerankComparison
from questionretrieval.retrievers.rerank.rerank_only_comparison import RerankOnlyComparison
from questionretrieval.retrievers.rerank.simplified_comparison import SimplifiedComparison

# 配置loguru日志
logger.remove()  # 移除默认handler
logger.add(
    sys.stderr,
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    level="INFO"
)


def load_questions_from_json(file_path: str) -> List[str]:
    """
    从JSON文件中加载问题列表
    
    Args:
        file_path: JSON文件路径
        
    Returns:
        List[str]: 问题列表
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            questions = json.load(f)
        return questions
    except Exception as e:
        logger.error(f"加载问题文件失败: {e}")
        raise


def prepare_sample_data() -> tuple[List[str], List[str]]:
    """
    准备示例数据
    
    Returns:
        tuple[List[str], List[str]]: (查询列表, 文档列表)
    """
    # 示例查询 - 基于水库调度文档设计
    queries = [
        "大伙房水库的总库容是多少？",
        "大伙房水库的主要供水对象有哪些？",
        "大伙房水库的防洪调度原则是什么？",
        "大伙房水库主汛期的防洪限制水位是多少？",
        "大伙房水库的调度责任部门包括哪些？"
    ]
    
    # 从JSON文件加载文档
    questions_file = os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        'questionretrieval', 'io', 'questions.json'
    )
    documents = load_questions_from_json(questions_file)
    
    return queries, documents


def run_embedding_rerank_experiment(queries: List[str], documents: List[str], 
                                   output_dir: str = "reports/rerank/embedding_rerank") -> EmbeddingRerankComparison:
    """
    运行embedding+rerank对比实验
    
    Args:
        queries: 查询列表
        documents: 文档列表
        output_dir: 输出目录
        
    Returns:
        EmbeddingRerankComparison: 实验对象
    """
    logger.info("开始embedding+rerank对比实验")
    
    # 初始化实验
    experiment = EmbeddingRerankComparison(
        embedding_model='BAAI/bge-m3',
        top_k_embedding=50,  # embedding阶段返回50个候选
        top_k_rerank=10,     # rerank阶段返回10个结果
        device=None,         # 自动检测设备
        qwen_config={
            'api_url': 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text-embedding/text-embedding',
            'api_key': 'your_qwen_api_key_here',  # 需要替换为实际的Qwen API密钥
            'timeout': 30
        }
    )
    
    # 加载embedding模型和文档
    experiment.load_embedding_model()
    experiment.load_documents(documents)
    
    # 运行批量对比
    results = experiment.run_batch_comparison(queries)
    
    # 保存结果
    experiment.save_results(output_dir)
    
    # 生成可视化
    experiment.visualize_results(output_dir)
    
    logger.info(f"embedding+rerank实验完成，结果保存在: {output_dir}")
    return experiment


def run_rerank_only_experiment(queries: List[str], documents: List[str], 
                              output_dir: str = "reports/rerank/rerank_only") -> RerankOnlyComparison:
    """
    运行rerank only对比实验
    
    Args:
        queries: 查询列表
        documents: 文档列表
        output_dir: 输出目录
        
    Returns:
        RerankOnlyComparison: 实验对象
    """
    logger.info("开始rerank only对比实验")
    
    # 初始化实验
    experiment = RerankOnlyComparison(
        top_k=10,
        device=None,  # 自动检测设备
        enable_parallel=True,
        qwen_config={
            'api_url': 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text-embedding/text-embedding',
            'api_key': 'your_qwen_api_key_here',
            'timeout': 30
        }
    )
    
    # 预加载所有模型
    experiment.load_all_models()
    
    # 运行批量对比（使用统一的候选集）
    results = experiment.run_uniform_candidates_comparison(queries, documents)
    
    # 保存结果
    experiment.save_results(output_dir)
    
    # 生成可视化
    experiment.visualize_results(output_dir)
    
    logger.info(f"rerank only实验完成，结果保存在: {output_dir}")
    return experiment


def run_comprehensive_analysis(embedding_experiment: EmbeddingRerankComparison,
                             rerank_experiment: RerankOnlyComparison,
                             output_dir: str = "reports/rerank/comprehensive") -> None:
    """
    运行综合分析
    
    Args:
        embedding_experiment: embedding+rerank实验对象
        rerank_experiment: rerank only实验对象
        output_dir: 输出目录
    """
    logger.info("开始综合分析")
    
    # 初始化简化分析
    comprehensive = SimplifiedComparison(output_dir)
    
    # 添加实验结果
    comprehensive.add_embedding_rerank_results(embedding_experiment.results)
    comprehensive.add_rerank_only_results(rerank_experiment.results)
    
    # 生成分析报告
    report_path = comprehensive.save_analysis_report()
    
    # 创建可视化
    comprehensive.create_visualization()
    
    logger.info(f"综合分析完成，结果保存在: {output_dir}")
    logger.info(f"综合报告: {report_path}")


def main():
    """
    主函数：运行完整的对比实验流程
    """
    logger.info("开始rerank模型对比实验")
    
    # 准备数据
    queries, documents = prepare_sample_data()
    logger.info(f"准备了 {len(queries)} 个查询和 {len(documents)} 个文档")
    
    # 创建输出目录
    base_output_dir = "reports/rerank"
    os.makedirs(base_output_dir, exist_ok=True)
    
    try:
        # 1. 运行embedding+rerank实验
        embedding_experiment = run_embedding_rerank_experiment(
            queries, documents, 
            os.path.join(base_output_dir, "embedding_rerank")
        )
        
        # 2. 运行rerank only实验
        rerank_experiment = run_rerank_only_experiment(
            queries, documents,
            os.path.join(base_output_dir, "rerank_only")
        )
        
        # 3. 运行综合分析
        run_comprehensive_analysis(
            embedding_experiment, rerank_experiment,
            os.path.join(base_output_dir, "comprehensive")
        )
        
        logger.info("所有实验完成！")
        logger.info(f"结果保存在: {base_output_dir}")
        
    except Exception as e:
        logger.error(f"实验过程中发生错误: {e}")
        raise


if __name__ == "__main__":
    main()